#!/usr/bin/python

import sys
import matplotlib.pyplot as plt
import pylab
import re
import random
from Bio import SeqIO
from Bio.SeqUtils import GC

"""

Usage: 
dissertation_GloeoAsmVerification.py glbkl_vs_AtleastOneAndSingletons.coords glbkl_vs_non-gloeo.coords glbkl_vs_454scaffolds.coords glbkl_vs_celeractgs.coords binned.gloeo.pairs.txt ../annotation/fixed.final_assembly_noCN4.fasta

Run the above command in this folder: /host/Users/JS/UH-work/gloeobacter/final_work/assembly_verification/

Usage: python gloeoAssemblyVerificationFigure2.py 1 2 3 4 5 
1: mummer coordinate file of assembled contig/genome vs. phymmBL-binned mate pairs and singletons
2: mummer coordinate file of assembled contig/genome vs. phymmBL-binnned non-gloeobacter reads
3: mummer coordinate file of assembled contig/genome vs. Newbler contigs
4: mummer coordinate file of assembled contig/genome vs. Contigs from other assemblers (Velvet or Celera)
5: mummer coordinate file of assembled contig/genome vs. mate pairs list (binned gloeobacter reads)
6: Fasta file of assembled contig(only one) or assembled genome

mate pair list looks like this:
GM6SIKE01AULG1  L->R    1 -> 9218       9093
GM6SIKE01B095V  L->R    1080 -> 11672   10261
GM6SIKE01ARQG1  L->R    2009 -> 13492   11328
GM6SIKE01AW9LY  L->R    3003 -> 13365   10120
GM6SIKE01BWLSW  L->R    4031 -> 14075   9896
GM6SIKE01BNO9P  R->L    5081 -> 13130   7966
GM6SIKE01CF48K  L->R    6078 -> 16829   10665
GM6SIKE01AWV4O  L->R    7036 -> 16355   9232

"""

##Functions
def fmt(f):
    st = '{0:.4}'.format(f)
    return st

def calAvg(countlist):
    total = 0
    count = len(countlist)
    for i in countlist:
        total += i
    avg = float(total/count)
    return avg

##regular expressions
m0 = re.compile('.*0$')
m1 = re.compile('.*1$')
m2 = re.compile('.*2$')
m3 = re.compile('.*3$')
m4 = re.compile('.*4$')
m5 = re.compile('.*5$')
m6 = re.compile('.*6$')
m7 = re.compile('.*7$')
m8 = re.compile('.*8$')
m9 = re.compile('.*9$')
left = re.compile('\w+_L')
right = re.compile('\w+_R')
#ctg = re.compile('sctg_\d+_\d+')

##Gloeobacter-specific reads
file1 = sys.argv[1]
f1 = open(file1, "rU")
fl1 = f1.readlines()

##Other reads
file2 = sys.argv[2]
f2 = open(file2, "rU")
fl2 = f2.readlines()

##454 contigs alignment
file3 = sys.argv[3]
f3 = open(file3, "rU")
fl3 = f3.readlines()

##Solexa contigs alignment
file4 = sys.argv[4]
f4 = open(file4, "rU")
fl4 = f4.readlines()

##phymmBL binning results

#Gloeobacter clone pairs
file5 = sys.argv[5]
f5 = open(file5, "rU")
fl5 = f5.readlines()

##Sequence file
seqfile = sys.argv[6]
sf = SeqIO.read(seqfile, "fasta")
gc_values = []
gcx = []

i = 0

while i < len(sf.seq):
    gc = GC(sf.seq[i:i+1000])
    gc_values.append(gc)
    gcx.append(i)
    i += 1000

##Parsing the file1 contents
line = fl1[4]
l = line.split('\t')
genome_size = int(l[7])

f1_coords = {}

for i in range(1, genome_size+1):
    f1_coords[i] = 0

pairsx1 = []
pairsx2 = []
pairsy = []

singletonsx1 = []
singletonsx2 = []
singletonsy = []

for i, line in enumerate(fl1[4:]):
    l = line.split('\t')
    start = int(l[0])
    stop = int(l[1])
    readname = l[12].rstrip()
    identity = float(l[6])
    for a in range(start, stop):
        f1_coords[a] = f1_coords[a] + 1

    if left.match(readname) or right.match(readname):
        pairsx1.append(start)
        pairsx2.append(stop)
        if m0.match(str(start)):
            pairsy.append(11.2)
        elif m1.match(str(start)):
            pairsy.append(11.4)
        elif m2.match(str(start)):
            pairsy.append(11.6)
        elif m3.match(str(start)):
            pairsy.append(11.8)
        elif m4.match(str(start)):
            pairsy.append(12)
        elif m5.match(str(start)):
            pairsy.append(12.2)
        elif m6.match(str(start)):
            pairsy.append(12.4)
        elif m7.match(str(start)):
            pairsy.append(12.6)
        elif m8.match(str(start)):
            pairsy.append(12.8)
        elif m9.match(str(start)):
            pairsy.append(13)
    else:
        singletonsx1.append(start)
        singletonsx2.append(stop)
        if m0.match(str(start)):
            singletonsy.append(11.2)
        elif m1.match(str(start)):
            singletonsy.append(11.4)
        elif m2.match(str(start)):
            singletonsy.append(11.6)
        elif m3.match(str(start)):
            singletonsy.append(11.8)
        elif m4.match(str(start)):
            singletonsy.append(12)
        elif m5.match(str(start)):
            singletonsy.append(12.2)
        elif m6.match(str(start)):
            singletonsy.append(12.4)
        elif m7.match(str(start)):
            singletonsy.append(12.6)
        elif m8.match(str(start)):
            singletonsy.append(12.8)
        elif m9.match(str(start)):
            singletonsy.append(13)

sxa = [singletonsx1, singletonsx2]
sxb = [singletonsy, singletonsy]

#sxa = [pairsx1, pairsx2]
#sxb = [pairsy, pairsy]

cov1 = []

for k, v in f1_coords.iteritems():
    cov1.append(v)

w = 1000
depthx1 = []
depthy1 = []
while w < len(cov1):
    average = calAvg(cov1[w-1000:w])
    depthx1.append(w-1000+500)
    depthy1.append(average)
    w += 1000

##Parsing the file2 contents

f2_coords = {}
for i in range(1, genome_size+1):
    f2_coords[i] = 0

for j, line in enumerate(fl2[4:]):
    l = line.split('\t')
    start = int(l[0])
    stop = int(l[1])
    identity = float(l[6])
    for a in range(start, stop):
        f2_coords[a] = f2_coords[a] + 1

cov2 = []

for k, v in f2_coords.iteritems():
    cov2.append(v)

x = 1000
depthx2 = []
depthy2 = []
while x < len(cov2):
    average = calAvg(cov2[x-1000:x])
    depthx2.append(x+500)
    depthy2.append(average)
    x += 1000

##For line representing the assembled genome
gx = [1,genome_size]
gy = [1,1]

xcoords = range(1, genome_size+1)
#ycoords = cov

##Parsing the file3 contents
##For 454 contigs assembled with Newbler
x1 = []
x2 = []
y454 = []

for i, line in enumerate(fl3[4:]):
    l = line.split('\t')
    start = int(l[0])
    stop = int(l[1])
    x1.append(start)
    x2.append(stop)
    if (i+1-1)%2 == 0:
        y454.append(6)
    else:
        y454.append(5.5)
a = [x1, x2]
b = [y454, y454]

##Parsing the file4 contents
##For Solex contigs assembled with Velvet
x3 = []
x4 = []
ys = []

for i, line in enumerate(fl4[4:]):
    l = line.split('\t')
    start = int(l[0])
    stop = int(l[1])
    x3.append(start)
    x4.append(stop)
    if (i+1-1)%2 == 0:
        ys.append(4)
    else:
        ys.append(3.5)
c = [x3, x4]
d = [ys, ys]

##Parsing the file5 contents
##Gloeobacter clone pairs
gpx1 = []
gpx2 = []
gpy = []

for i, line in enumerate(fl5):
    l = line.split('\t')
    start = int(l[2].split(' ')[0])
    stop = int(l[2].split(' ')[2])
    gpx1.append(start)
    gpx2.append(stop)
    if m0.match(str(i)):
        gpy.append(7.2)
    elif m1.match(str(i)):
        gpy.append(7.4)
    elif m2.match(str(i)):
        gpy.append(7.6)
    elif m3.match(str(i)):
        gpy.append(7.8)
    elif m4.match(str(i)):
        gpy.append(8)
    elif m5.match(str(i)):
        gpy.append(8.2)
    elif m6.match(str(i)):
        gpy.append(8.4)
    elif m7.match(str(i)):
        gpy.append(8.6)
    elif m8.match(str(i)):
        gpy.append(8.8)
    elif m9.match(str(i)):
        gpy.append(9)

gpa = [gpx1, gpx2]
gpb = [gpy, gpy]

##Close the file handlers
f1.close()
f2.close()
f3.close()
f4.close()

##Print notice..
print "# of records in xcoords: ", len(xcoords)
print "# of records in cov1: ", len(cov1)

##Start plotting
fig = plt.figure(1, figsize=(15,8))
#plt.subplots_adjust(wspace=4.0)

##First subplot
#plt.subplot(311)
#plt.axis([0, genome_size, 0, 20])
#plt.plot(gpa, gpb, color='purple', linestyle='-')
#plt.title('Gloeobacter assembly verification')
#plt.grid(True)

##First subplot
ax1 = fig.add_subplot(211)
#plt.subplot(311)
#plt.title('Assembly verification')
ax1.plot(gx, gy, color='#AF7817', marker='|', markersize=8.0,
    mec='black', ls='-', lw=2.0)
ax1.plot(sxa, sxb, color='#333366', linestyle='-')
ax1.plot(a, b, color='purple', linestyle='-', lw=2.0)
ax1.plot(c, d, color='red', linestyle='-', lw=2.0)
ax1.plot(gpa, gpb, color='#ADA96E', linestyle='-')
ax1.annotate('Gloeobacter singleton reads', xy=(0.8, 0.95), 
    xycoords='axes fraction', horizontalalignment='center', 
    verticalalignment='center', fontsize=9)
ax1.annotate('Mate pairs (at least one is binned as Gloeobacter)', 
    xy=(0.8, 0.68), xycoords='axes fraction', horizontalalignment='center', 
    verticalalignment='center', fontsize=9)
ax1.annotate('Newbler contigs (454)', xy=(0.8, 0.47), xycoords='axes fraction', 
    horizontalalignment='center', verticalalignment='center', fontsize=9)
ax1.annotate('Celera contigs (454+Illumina)', xy=(0.8, 0.20), 
    xycoords='axes fraction', horizontalalignment='center', 
    verticalalignment='center', fontsize=9)

ax1.axis([0, genome_size, 0, 14])
#ax1.axis([415517, 431636, 0, 14]) #region1
#ax1.axis([903207, 946727, 0, 14]) #region2
#ax1.axis([1004060, 1015740, 0, 14]) #region3
#ax1.axis([1732580, 1828970, 0, 14]) #region4
#ax1.axis([4262380, 4279810, 0, 14]) #region5
#ax1.axis([860000, 980000, 0, 14]) #region 1
#ax1.axis([1600000, 2100000, 0, 14]) #region 2
#ax1.axis([4200000, 4550000, 0, 14]) #region 3

frame1 = plt.gca()
for tick in frame1.axes.get_yticklines():
    tick.set_visible(False)
for y in frame1.axes.get_yticklabels():
    y.set_visible(False)
ax1.grid(False)

##Second subplot
ax2 = fig.add_subplot(212)
#plt.subplot(312)
ax2.fill_between(depthx2, depthy2, facecolor='#33FF33', alpha=0.4)
ax2.fill_between(depthx1, depthy1, facecolor='#660066', alpha=0.4)
#ax2.annotate('Green = other organisms', xy=(0.2, 0.80), 
#    xycoords='axes fraction', horizontalalignment='center', 
#    verticalalignment='center')
#ax2.annotate('Purple = Gloeobacter', xy=(0.2, 0.60), xycoords='axes fraction', 
#    horizontalalignment='center', verticalalignment='center')
ax2.axvspan(2728368, 2733167, facecolor='blue', alpha=0.4)
ax2.annotate('rRNA operon', xy=(2728368, 40), xycoords='data', xytext=(30, 0), 
    textcoords='offset points', arrowprops=dict(arrowstyle="->"))
ax2.axis([0, genome_size, 0, max(depthy1)+2])
#ax2.axis([415517, 431636, 0, max(depthy1)+2])
#ax2.axis([903207, 946727, 0, max(depthy1)+2])
#ax2.axis([1004060, 1015740, 0, max(depthy1)+2])
#ax2.axis([1732580, 1828970, 0, max(depthy1)+2])
#ax2.axis([4262380, 4279810, 0, max(depthy1)+2])
#ax2.axis([860000, 980000, 0, max(depthy1)+2]) #region 1
#ax2.axis([1600000, 2100000, 0, max(depthy1)+2]) #region 2
#ax2.axis([4200000, 4550000, 0, max(depthy1)+2]) #region 3
ax2.set_xlabel('Genome position (bp)')
ax2.set_ylabel('Counts of binned reads (per 1000bp)')

ax3 = ax2.twinx()
ax3.plot(gcx, gc_values, color='blue', linestyle='-', alpha=0.4)
ax3.axis([0, genome_size, 0, 80])
#ax3.axis([415517, 431636, 0, 80])
#ax3.axis([903207, 946727, 0, 80])
#ax3.axis([1004060, 1015740, 0, 80])
#ax3.axis([1732580, 1828970, 0, 80])
#ax3.axis([4262380, 4279810, 0, 80])

#ax3.axis([860000, 980000, 0, 80]) #region 1
#ax3.axis([1600000, 2100000, 0, 80]) #region 2
#ax3.axis([4200000, 4550000, 0, 80]) #region 3

ax3.set_ylabel('G+C %')
#for tl in ax3.get_yticklabels():
#    tl.set_color('b')

#ax4 = ax2.twinx()
#ax4.plot(gpa, gpb, color='#ADA96E', linestyle='-')
#ax4.axis([0, genome_size, 0, 10])

#frame1 = plt.gca()
#for tick in frame1.axes.get_yticklines():
#    tick.set_visible(False)
#for y in frame1.axes.get_yticklabels():
#    y.set_visible(False)

ax1.axvspan(415517, 431636, facecolor='#FFD700', alpha=0.2)
ax1.axvspan(903207, 946727, facecolor='#FFD700', alpha=0.2)
ax1.axvspan(1004060, 1015740, facecolor='#FFD700', alpha=0.2)
ax1.axvspan(1732580, 1828970, facecolor='#FFD700', alpha=0.2)
ax1.axvspan(4262380, 4279810, facecolor='#FFD700', alpha=0.2)

ax2.axvspan(415517, 431636, facecolor='#FFD700', alpha=0.2)
ax2.axvspan(903207, 946727, facecolor='#FFD700', alpha=0.2)
ax2.axvspan(1004060, 1015740, facecolor='#FFD700', alpha=0.2)
ax2.axvspan(1732580, 1828970, facecolor='#FFD700', alpha=0.2)
ax2.axvspan(4262380, 4279810, facecolor='#FFD700', alpha=0.2)

ax2.grid(True)

##Third subplot

#plt.subplot(313)
#plt.plot(gcx, gc_values, color='blue', linestyle='-', alpha=0.4)
#plt.xlabel("Genes")
#plt.ylabel("G+C%")

plt.show()

#plt.savefig(outfile, format='pdf')
